package algorithm.permutation;

import java.util.Arrays;
import java.util.function.Consumer;
import java.util.stream.IntStream;

public class MagicSquare {
    int magicSuqareLength;
    int squareCount;
    int[] dataArray;
    int lineSum;

    public MagicSquare(int magicSuqareLength) {
        this.magicSuqareLength = magicSuqareLength;
        this.squareCount = magicSuqareLength * magicSuqareLength;
        this.dataArray = IntStream.range(1, squareCount + 1).toArray();
        lineSum = Arrays.stream(this.dataArray).sum() / this.magicSuqareLength;
    }

    public void make(Consumer<int[]> onSuccess) {
        this.visit(it -> {
            if (check(it)) {
                onSuccess.accept(it);
            }
        });
    }

    private boolean check(int[] square) {
        for (int i = 0; i < magicSuqareLength; i++) {
            int sum = 0;
            for (int j = 0; j < magicSuqareLength; j++) {
                sum += square[i * magicSuqareLength + j];
            }
            if (sum != lineSum) {
                return false;
            }
        }
        for (int i = 0; i < magicSuqareLength; i++) {
            int sum = 0;
            for (int j = 0; j < magicSuqareLength; j++) {
                sum += square[j * magicSuqareLength + i];
            }
            if (sum != lineSum) {
                return false;
            }
        }
        {
            int sum = 0;
            for (int i = 0; i < magicSuqareLength; i++) {
                sum += square[i * magicSuqareLength + i];
            }
            if (sum != lineSum) {
                return false;
            }
        }
        {
            int sum = 0;
            for (int i = 0; i < magicSuqareLength; i++) {
                sum += square[i * magicSuqareLength + magicSuqareLength - i - 1];
            }
            if (sum != lineSum) {
                return false;
            }
        }
        return true;
    }

    public void visit(Consumer<int[]> onVisit) {
        visit(dataArray.length, onVisit);
    }

    public void visit(int length, Consumer<int[]> onVisit) {
        int[] visitedState = new int[dataArray.length];
        Arrays.fill(visitedState, 1);
        int[] newDataArray = Arrays.copyOf(dataArray, dataArray.length);
        doVisit(newDataArray, length, onVisit, new int[dataArray.length], 0, visitedState);
    }

    private void doVisit(int[] dataArray, int length, Consumer<int[]> onVisit, int[] context, int validContextIndex, int[] visitedState) {
        if (validContextIndex == length) {
            onVisit.accept(Arrays.copyOf(context, length));
        } else {
            for (int i = 0; i < dataArray.length; i++) {
                if (visitedState[i] > 0) {
                    visitedState[i]--;
                    context[validContextIndex] = dataArray[i];
                    if (validContextIndex != 0 && validContextIndex % this.magicSuqareLength == 0) {
                        if (lineCheck(context, validContextIndex / this.magicSuqareLength - 1)) {
                            doVisit(dataArray, length, onVisit, context, validContextIndex + 1, visitedState);
                        }
                    } else {
                        doVisit(dataArray, length, onVisit, context, validContextIndex + 1, visitedState);
                    }
                    visitedState[i]++;
                }
            }
        }
    }

    private boolean lineCheck(int[] context, int line) {
        int sum = 0;
        for (int j = 0; j < magicSuqareLength; j++) {
            sum += context[line * magicSuqareLength + j];
        }
        if (sum != lineSum) {
            return false;
        }
        return true;
    }
}
